# Copyright (c) Facebook, Inc. and its affiliates.
import logging
from copy import deepcopy
from typing import Callable, Dict, List, Optional, Tuple, Union

import fvcore.nn.weight_init as weight_init

import torch
from torch import nn
from torch.nn import functional as F

import numpy as np

from detectron2.config import configurable
from detectron2.layers import Conv2d, ShapeSpec, get_norm
from detectron2.modeling import SEM_SEG_HEADS_REGISTRY
from detectron2.structures import Boxes, ImageList, Instances, BitMasks
from detectron2.modeling.postprocessing import sem_seg_postprocess
from detectron2.utils.memory import retry_if_cuda_oom

from ..transformer_decoder.maskformer_transformer_decoder import build_transformer_decoder
from ..pixel_decoder.fpn import build_pixel_decoder
from ..retrieval_decoder import DefaultInstanceRetrievalDecoder
from ..retrieval_decoder.roi_enhance import ROIEnhance

@SEM_SEG_HEADS_REGISTRY.register()
class BuildingRetrievalHeadDefault(nn.Module):
    _version = 2

    def _load_from_state_dict(
        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
    ):
        version = local_metadata.get("version", None)
        if version is None or version < 2:
            # Do not warn if train from scratch
            scratch = True
            logger = logging.getLogger(__name__)
            for k in list(state_dict.keys()):
                newk = k
                if "sem_seg_head" in k and not k.startswith(prefix + "predictor"):
                    newk = k.replace(prefix, prefix + "pixel_decoder.")
                    # logger.debug(f"{k} ==> {newk}")
                if newk != k:
                    state_dict[newk] = state_dict[k]
                    del state_dict[k]
                    scratch = False

            if not scratch:
                logger.warning(
                    f"Weight format of {self.__class__.__name__} have changed! "
                    "Please upgrade your models. Applying automatic conversion now ..."
                )

    @configurable
    def __init__(
        self,
        input_shape: Dict[str, ShapeSpec],
        *,
        num_classes: int,
        pixel_decoder: nn.Module,
        loss_weight: float = 1.0,
        ignore_value: int = -1,
        # extra parameters
        transformer_predictor: nn.Module,
        transformer_in_feature: str,
        train_retrieval_on: bool = True,
        train_mask_on: bool = True,
        roi_expand:bool = False,
        roi_shift:bool = False
    ):
        """
        NOTE: this interface is experimental.
        Args:
            input_shape: shapes (channels and stride) of the input features
            num_classes: number of classes to predict
            pixel_decoder: the pixel decoder module
            loss_weight: loss weight
            ignore_value: category id to be ignored during training.
            transformer_predictor: the transformer decoder that makes prediction
            transformer_in_feature: input feature name to the transformer_predictor
        """
        super().__init__()
        input_shape = sorted(input_shape.items(), key=lambda x: x[1].stride)
        self.in_features = [k for k, v in input_shape]
        feature_strides = [v.stride for k, v in input_shape]
        feature_channels = [v.channels for k, v in input_shape]

        self.ignore_value = ignore_value
        self.common_stride = 4
        self.loss_weight = loss_weight

        self.pixel_decoder = pixel_decoder

        #decoder
        self.predictor = transformer_predictor
        self.transformer_in_feature = transformer_in_feature        
        
        self.train_mask_on = train_mask_on          
        self.train_retrieval_on = train_retrieval_on
        if self.train_retrieval_on:
            #instance decoder
            self.instance_decoder = DefaultInstanceRetrievalDecoder(input_channel=feature_channels[0])
            self.roi_enhance = ROIEnhance(shift=roi_shift,expand=roi_expand)
        self.num_classes = num_classes

        
    @classmethod
    def from_config(cls, cfg, input_shape: Dict[str, ShapeSpec]):
        # figure out in_channels to transformer predictor
        if cfg.MODEL.MASK_FORMER.TRANSFORMER_IN_FEATURE == "transformer_encoder":
            transformer_predictor_in_channels = cfg.MODEL.SEM_SEG_HEAD.CONVS_DIM
        elif cfg.MODEL.MASK_FORMER.TRANSFORMER_IN_FEATURE == "pixel_embedding":
            transformer_predictor_in_channels = cfg.MODEL.SEM_SEG_HEAD.MASK_DIM
        elif cfg.MODEL.MASK_FORMER.TRANSFORMER_IN_FEATURE == "multi_scale_pixel_decoder":  # for maskformer2
            transformer_predictor_in_channels = cfg.MODEL.SEM_SEG_HEAD.CONVS_DIM
        else:
            transformer_predictor_in_channels = input_shape[cfg.MODEL.MASK_FORMER.TRANSFORMER_IN_FEATURE].channels

        return {
            "input_shape": {
                k: v for k, v in input_shape.items() if k in cfg.MODEL.SEM_SEG_HEAD.IN_FEATURES
            },
            "ignore_value": cfg.MODEL.SEM_SEG_HEAD.IGNORE_VALUE,
            "num_classes": cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES,
            "pixel_decoder": build_pixel_decoder(cfg, input_shape),
            "loss_weight": cfg.MODEL.SEM_SEG_HEAD.LOSS_WEIGHT,
            "transformer_in_feature": cfg.MODEL.MASK_FORMER.TRANSFORMER_IN_FEATURE,
            "transformer_predictor": None,
            "train_retrieval_on": cfg.MODEL.SEM_SEG_HEAD.TRAIN_RETRIEVAL_ON,
            "train_mask_on": cfg.MODEL.SEM_SEG_HEAD.TRAIN_MASK_ON,
            "roi_shift": cfg.MODEL.SEM_SEG_HEAD.ROI_SHIFT,
            "roi_expand": cfg.MODEL.SEM_SEG_HEAD.ROI_EXPAND
        }

    def forward(self, features, targets=None):
        #mask no use
        mask = None
        #mask_features encoder的最后一次层输出
        #multi_scale_features encoder的多层输出 3层
        if self.pixel_decoder is not None:
            mask_features, transformer_encoder_features, multi_scale_features = self.pixel_decoder.forward_features(features)
            # mask_features = multi_scale_features[-1]
        else:
            mask_features = features[self.in_features[0]]
        predictions = {}
        
        if self.train_retrieval_on:
            boxes_tensor,b_tensor,ids_tensor = self.roi_enhance.getBoxesinTargets(targets,scale=4)
            perd_instances = self.instance_decoder(mask_features,boxes_tensor,b_tensor)
            predictions['instance_perd_feats'] = perd_instances
            predictions['instance_tgt_ids'] = ids_tensor            
        return predictions
    
    def forward_eval_retrieval(self,features,batched_inputs, targets=None):
        processed_results = []
        bs = len(targets)
        if self.pixel_decoder is not None:
            mask_features, transformer_encoder_features, multi_scale_features = self.pixel_decoder.forward_features(features)
            # mask_features = multi_scale_features[-1]
        else:
            mask_features = features[self.in_features[0]]
        image_size = targets[0]["masks"].shape[-2:]
        ori_height = batched_inputs[0].get("height", image_size[0])
        ori_width = batched_inputs[0].get("width", image_size[1])
        for b in range(bs):
            target = targets[b]
            boxes,valid_boxes,valid_box_index = self.roi_enhance.getBoxesofTargetCFG(target)
            # tgt_masks = targets[b]["masks"]
            tgt_ids = target["gt_ids"]
            # boxes = BitMasks(tgt_masks).get_bounding_boxes()
            # valid_box_index = boxes.nonempty(threshold=16.0)
            # valid_boxes = boxes[valid_box_index]
            if(len(valid_boxes) < 1):
                continue
            valid_features = self.instance_decoder.inference_one(mask_features[[b]],valid_boxes)
            pred_features = torch.zeros((valid_box_index.shape[0],valid_features.shape[-1])).to(valid_features)
            pred_features[valid_box_index] = valid_features
            
            result = Instances([ori_height,ori_width])
            boxes.tensor = boxes.tensor * ori_height / image_size[0]
            result.pred_boxes = boxes
            result.pred_boxes_ins_mask = valid_box_index
            result.pred_boxes_ins_feature = pred_features
            result.tgt_ins_ids = tgt_ids
            result.scores = torch.ones((boxes.tensor.shape[0]))
            result.pred_classes = torch.zeros((boxes.tensor.shape[0]),dtype=torch.int)
            
            processed_results.append({})
            processed_results[-1]["instances"] = result
        return processed_results
#region
    def forward_eval(self,features,batched_inputs, images, targets=None):
        #
        mask = None
        mask_features, transformer_encoder_features, multi_scale_features = self.pixel_decoder.forward_features(features)
        if self.transformer_in_feature == "multi_scale_pixel_decoder":
            predictions = self.predictor(multi_scale_features, mask_features, mask)
        else:
            if self.transformer_in_feature == "transformer_encoder":
                assert (
                    transformer_encoder_features is not None
                ), "Please use the TransformerEncoderPixelDecoder."
                predictions = self.predictor(transformer_encoder_features, mask_features, mask)
            elif self.transformer_in_feature == "pixel_embedding":
                predictions = self.predictor(mask_features, mask_features, mask)
            else:
                predictions = self.predictor(features[self.transformer_in_feature], mask_features, mask)        
        
        mask_cls_results = predictions["pred_logits"]
        mask_pred_results = predictions["pred_masks"]
        # upsample masks
        mask_pred_results = F.interpolate(
            mask_pred_results,
            size=(images.tensor.shape[-2], images.tensor.shape[-1]),
            mode="bilinear",
            align_corners=False,
        )
        
        processed_results = []
        b_tensor = torch.tensor([],dtype=torch.int)
        boxes_tensor = torch.tensor([],dtype=torch.int)
        
        batch_index = -1
        for mask_cls_result, mask_pred_result, input_per_image, image_size in zip(
            mask_cls_results, mask_pred_results, batched_inputs, images.image_sizes
        ):
            batch_index += 1
            height = input_per_image.get("height", image_size[0])
            width = input_per_image.get("width", image_size[1])

            processed_results.append({})


            if self.train_retrieval_on:
                ### 计算INSTANCE MIX FEATURE
                pred_boxes = BitMasks(mask_pred_result > 0).get_bounding_boxes()

                ###
            # 这个函数的功能是将mask_pred_result 调整为原始图像输入的H和W
            mask_pred_result = retry_if_cuda_oom(sem_seg_postprocess)(
                mask_pred_result, image_size, height, width
            )
            instance_r,resort_topk_indices = retry_if_cuda_oom(self.instance_mask_inference)(mask_cls_result, mask_pred_result)
            
            if self.train_retrieval_on:
                resort_topk_indices = resort_topk_indices.to(pred_boxes.tensor.device)
                pred_boxes.tensor = pred_boxes.tensor[resort_topk_indices]
                valid_box_index = pred_boxes.nonempty(threshold=8.0)
                valid_boxes = pred_boxes[valid_box_index]
                pred_feature = torch.zeros((pred_boxes.tensor.shape[0],self.instance_decoder.out_feature_dim))
                
                if(len(valid_boxes) > 0):
                    valid_boxes = np.ceil(valid_boxes.tensor / 4 )
                    valid_boxes= valid_boxes.to(torch.int)
                    perd_mix_feats = self.instance_decoder.inference(mask_features,multi_scale_features,batch_index,valid_boxes)
                    # pred_feature = pred_feature.to(perd_mix_feats)
                    perd_mix_feats = perd_mix_feats.to(pred_feature)
                    pred_feature[valid_box_index] = perd_mix_feats
                instance_r.pred_boxes_ins_mask = valid_box_index
                instance_r.pred_boxes_ins_feature = pred_feature
                
            processed_results[-1]["instances"] = instance_r
            
        return processed_results
   
    def instance_mask_inference(self, mask_cls, mask_pred):
        
        test_topk_per_image = 100
        num_queries = 100
        # mask_pred is already processed to have the same shape as original input
        image_size = mask_pred.shape[-2:]

        # [Q, K]
        scores = F.softmax(mask_cls, dim=-1)[:, :-1]
        labels = torch.arange(self.num_classes).unsqueeze(0).repeat(num_queries, 1).flatten(0, 1)
        labels = labels.to(scores)
        # scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.num_queries, sorted=False)
        scores_per_image, topk_indices = scores.flatten(0, 1).topk(test_topk_per_image, sorted=False)
        labels_per_image = labels[topk_indices]
        
        topk_indices = topk_indices // self.num_classes
        #mask_pred = mask_pred.unsqueeze(1).repeat(1, self.sem_seg_head.num_classes, 1).flatten(0, 1)
        mask_pred = mask_pred[topk_indices]

        result = Instances(image_size)
        # mask (before sigmoid)
        result.pred_masks = (mask_pred > 0).float()
        
        #result.pred_boxes = Boxes(torch.zeros(mask_pred.size(0), 4))
        # Uncomment the following to get boxes from masks (this is slow)
        result.pred_boxes = BitMasks(mask_pred > 0).get_bounding_boxes()

        # calculate average mask prob
        mask_scores_per_image = (mask_pred.sigmoid().flatten(1) * result.pred_masks.flatten(1)).sum(1) / (result.pred_masks.flatten(1).sum(1) + 1e-6)
        result.scores = scores_per_image * mask_scores_per_image
        result.pred_classes = labels_per_image.to(torch.int)
        return result,topk_indices
